#coding=utf-8
#这个程序主要用于识别医疗发票的名字、收费细项和总金额，先用opencv从图片中定位出关键区域，进而调用腾讯ocr的接口来识别内容
import cv2
import numpy as np
import os
import matplotlib.pyplot as plt
# from mst import tencent_ocr
import sys
from PIL import Image
import pytesseract

def matchWordTemplate(gray_im, Img_word):
    wordPos = []  # 满足条件字体位置
    res = cv2.matchTemplate(gray_im, Img_word, cv2.TM_CCOEFF_NORMED)
    threshold = 0.7
    loc = np.where(res >= threshold)
    # print loc
    # print zip(*loc[::-1])
    for pt in zip(*loc[::-1]):
        # if (pt[0] > int(gray_im.shape[0] / 3)):
        wordPos.append(pt)
            # cv2.rectangle(im, pt, (pt[0] + Img_word.shape[1], pt[1] + Img_word.shape[0]), (0, 0, 255), 2)
    if len(wordPos) == 0:
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)

        # if (max_loc > 0.5):# & (max_loc[0] > int(gray_im.shape[0] / 3)):  # 判断位置，只有高于图片1/4的位置被识别
        if max_val > 0.4:
            wordPos.append(max_loc)
            # cv2.rectangle(im, (max_loc[0], max_loc[1]), (max_loc[0] + Img_word.shape[1], max_loc[1] + Img_word.shape[0]), (255, 0, 0),2)  # 画出边界框（正的矩形）
    return wordPos


def getMatchWordTemplateSimilarity(gray_im, Img_word):
    # 获取匹配的相似度
    res = cv2.matchTemplate(gray_im, Img_word, cv2.TM_CCOEFF_NORMED)
    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
    return max_val


def getBestScale(gray_im):
    # 获取最佳放缩倍数
    H, W = gray_im.shape[:2]  # 获取图像高度和宽度
    scaleSet = np.linspace(0.35, 1.5, 8)  # 缩放集合
    scaleSimSet = []  # 缩放相似性集合
    for scale in scaleSet:  # np.arange(0.35,1.5,0.25):
        Img_scale_im = cv2.resize(gray_im, (int(W * scale), int(H * scale)))  # 重新放缩
        simSet = []  # 存储本次图片放缩的拟合度集合
        for png in os.listdir('templates'):
            Img_word = cv2.imread("templates/" + png, 0)  # 读取字体模板
            sim = getMatchWordTemplateSimilarity(Img_scale_im, Img_word)
            if sim > 0.5:  # 取相似性大于0.5
                simSet.append(sim)
        # print(scale, simSet)
        if (len(simSet) > 0):
            scaleSimSet.append(np.mean(simSet))
        else:
            scaleSimSet.append(0)
    # print scaleSimSet
    scale = scaleSet[np.argmax(scaleSimSet)]  # 获取最佳放缩倍数
    return H, W, scale


def getKeyWordInfo(gray_im, im, flag=1):
    # 获取关键字眼信息
    # keyWordInfo:numpy.array,size=(,4),存储所有关键字的位置信息，[x,y,delta_x,delta_y]
    # x:关键字左上角宽度坐标，y关键字左上角高度坐标，delta_x,关键字宽度，delta_y，关键字高度

    # keyWordInfoWordAndPosition，dict(),存储每个关键字的位置信息，key为关键字
    # 每个key里面内容为numpy.array,size=(,4),存储所有关键字的位置信息，[x,y,delta_x,delta_y]
    # x:关键字左上角宽度坐标，y关键字左上角高度坐标，delta_x,关键字宽度，delta_y，关键字高度
    keyWordInfo = []
    keyWordInfoWordAndPosition = dict()
    for png in os.listdir('templates'):
        Img_word = cv2.imread("templates/" + png, 0)  # 读取字模板
        # print("template2/" + png)
        wordPos = matchWordTemplate(gray_im, Img_word)
        keyWordInfoWordAndPosition[png] = []
        if len(wordPos) > 0:
            for pos in wordPos:
                keyWordInfo.append([pos[0], pos[1], Img_word.shape[1], Img_word.shape[0]])
                keyWordInfoWordAndPosition[png].append([pos[0], pos[1], Img_word.shape[1], Img_word.shape[0]])
    if len(keyWordInfo) > 0:
        if len(keyWordInfo) > 1:
            keyWordInfo = sorted(keyWordInfo, key=lambda x: (x[0], x[1]))
            # print keyWordInfo
            keyWordInfo = np.array(keyWordInfo)
            temp = np.diff(keyWordInfo, axis=0)
            judegeX = np.hstack([1000, np.abs(temp[:, 0])])
            judegeY = np.hstack([1000, np.abs(temp[:, 1])])
            keyWordInfo = keyWordInfo[(judegeX > 5) | (judegeY > 5), :]
        else:
            keyWordInfo = np.array(keyWordInfo)
        # if flag == 1:
        #     for i in range(keyWordInfo.shape[0]):
        #         cv2.rectangle(im, (keyWordInfo[i, 0], keyWordInfo[i, 1]),
        #                       (keyWordInfo[i, 0] + keyWordInfo[i, 2], keyWordInfo[i, 1] + keyWordInfo[i, 3]),
        #                       (0, 0, 255), 10)
    # 判断
    keyWordInfoRowNum = keyWordInfo.shape[0]
    for png in os.listdir('templates'):
        thisPngKeyWord = keyWordInfoWordAndPosition[png]
        thisPngKeyWord = np.array(thisPngKeyWord)
        getKeyWordList = []
        for i in range(thisPngKeyWord.shape[0]):
            ith_thisPngKeyWord = np.tile(thisPngKeyWord[i,:],(keyWordInfoRowNum ,1))
            getKeyWordList.append(np.any(np.sum(ith_thisPngKeyWord-keyWordInfo,axis=1)==0))
        keyWordInfoWordAndPosition[png] = thisPngKeyWord[getKeyWordList,:]
    return keyWordInfo,keyWordInfoWordAndPosition


def getWordArea(gray_im):
    # kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (7,7))
    kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (15, 15))
    gray = cv2.morphologyEx(gray_im, cv2.MORPH_OPEN, kernel)
    ret, thresh = cv2.threshold(gray, 127, 255, 0)
    thresh = 255 - thresh
    images, contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)  # 对每一个连贯的形状画出边框
    cSet = []
    for c in contours:
        x, y, w, h = cv2.boundingRect(c)
        cSet.append([x, y, w, h])
    cSet = sorted(cSet, key=lambda x: (x[2], x[3]), reverse=True)
    cSet = np.array(cSet)
    findElementArea = []
    curElement = cSet[0, :]
    cSet = cSet[1:, :]
    findElementArea.append(list(curElement))
    while len(cSet) > 0:
        judge1 = (cSet[:, 0] > curElement[0]) & (cSet[:, 1] > curElement[1])
        judge2 = ((cSet[:, 0] + cSet[:, 2]) <= (curElement[0] + curElement[2])) & (
                    (cSet[:, 1] + cSet[:, 3]) <= (curElement[1] + curElement[3]))
        judge = judge1 & judge2
        cSet = cSet[~judge, :]
        curElement = cSet[0, :]
        cSet = cSet[1:, :]
        findElementArea.append(list(curElement))
        # print(cSet.shape)
    # for area in findElementArea:
    #     cv2.rectangle(im, (area[0], area[1]), (area[0]+area[2], area[1]+area[3]), (255, 255, 0), 2)
    return findElementArea

def gamma(im, k): #gamma变换
    em=im.astype("float32")
    # height,width=im.shape[:2]
    fm=np.zeros(em.shape, dtype=np.float32)
    em=em**k
    for i in range(3):
        a = np.zeros(em.shape[:2], dtype=np.float32)
        cv2.normalize(em[:,:,i],a,0,255,cv2.NORM_MINMAX)
        fm[:,:,i]=a
    return fm.astype("uint8")


def find_retangles(gray, scale): #用于寻找图片中符合要求的边框
    gray = cv2.blur(gray, (3, 3))
    median = np.median(gray)
    yuzhi = 100
    tangles = 0
    save = 0
    while yuzhi < 200: #寻找二值化的最佳阈值
        ret, thresh = cv2.threshold(gray, yuzhi, 255, cv2.THRESH_BINARY)
        # thresh=cv2.medianBlur(thresh, 3)
        points = []
        image, contours, hier = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for c in contours:
            x, y, w, h = cv2.boundingRect(c)
            if w > h * 2 and w > 100*scale and h>50*scale:
                # cv2.rectangle(im, (x,y), (x+w,y+h), (255,0,0), 3)
                points.append([x, y, w, h])
        if len(points) > tangles:
            tangles = len(points)
            save = yuzhi #保存阈值
        yuzhi += 5

    # print(save)
    ret, thresh = cv2.threshold(gray, save, 255, cv2.THRESH_BINARY) #二值化
    points = []
    image, contours, hier = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for c in contours:
        x, y, w, h = cv2.boundingRect(c)
        if w > h * 2 and w > 100*scale and h > 50*scale: #设定边框选取条件
            # cv2.rectangle(im, (x,y), (x+w,y+h), (255,0,0), 2)
            points.append([x, y, w, h])
            # image=thresh[y:y+h, x:x+w]
    return points


if len(sys.argv)>1: #读取参数名作为文件名
    im = cv2.imread(sys.argv[1])
else:
    im=cv2.imread("shouju.jpg")
gray=im[:,:,0]

# print rectangles
# for rectangle in rectangles:
#     x,y,w,h=rectangle
#     cv2.rectangle(im, (x, y), (x + w, y + h), (255, 0, 0), 2)
H,W,scale=getBestScale(gray) #获得原图的最佳缩放倍数
gray=cv2.resize(gray, (int(W * scale), int(H * scale)))
im = cv2.resize(im, (int(W * scale), int(H * scale))) #缩放原图
fm = im.copy()

rectangles=np.array(find_retangles(gray, scale)) #获取边框
for rectangle in rectangles:
    x,y,w,h=rectangle
    # cv2.rectangle(im, (x, y), (x + w, y + h), (0, 0, 0), 5) ####
# copyone=rectangles.copy()
_, keyWordInfoWordAndPosition = getKeyWordInfo(gray, im) #找到关键字的每一个位置
# print(keyWordInfoWordAndPosition)


for key in keyWordInfoWordAndPosition.keys(): #同一个关键字可能在同一个地方重复的给出了两个位置信息，需要去除重复的位置信息
    temp=keyWordInfoWordAndPosition[key]
    # print temp
    indextodelete=[]
    uniquedata=[]
    if len(temp)>1:
        for i in range(len(temp)-1):
            for j in range(i+1, len(temp)): #使用双重循环，如果后一个和前一个重复了，则删除后一个
                if abs(temp[j][0]-temp[i][0])<5:
                    indextodelete.append(j)
    if len(indextodelete)>0:
        for i in range(len(temp)):
            if i not in indextodelete:
                uniquedata.append(temp[i]) #保存下不重复的位置
        keyWordInfoWordAndPosition[key]=uniquedata
# for key in keyWordInfoWordAndPosition:
#     temp=keyWordInfoWordAndPosition[key]
    # for word in temp:
    #     cv2.rectangle(im, (word[0],word[1]), (word[0]+word[2], word[1]+word[3]), (0,0,255), 5) ###

# print(keyWordInfoWordAndPosition)
# keywordrectangles={} #创建一个字典，字典的key为每一个关键字
# for key in keyWordInfoWordAndPosition.keys():
#     keywordrectangles[key]={} #创建双重字典，字典的key为每一个关键字的具体位置
#     for word in keyWordInfoWordAndPosition[key]:
#         keywordrectangles[key][word[0]]=[] # 字典的关键字的内容为一个列表，用于存放对应的蓝色矩形
# # print keywordrectangles
# for key in keyWordInfoWordAndPosition.keys():
#     temp=keyWordInfoWordAndPosition[key]
#     for word in temp:
#         xdistance=1000
#         absydistance=1000
#         xindex=None
#
#         ydistance=1000
#         absxdistance=1000
#         yindex=None
#
#         for i in range(len(rectangles)):
#             x,y,w,h=rectangles[i]
#             if x-word[0]-word[2]> 0 and np.sqrt((x+w/2.-word[0])**2+(y+h/2.-word[1])**2)<=xdistance and np.abs(y-word[1])<im.shape[0]*0.055: #50
#                 xdistance=np.sqrt((x+w/2.-word[0])**2+(y+h/2.-word[1])**2)
#                 # absydistance=abs(y-word[1])
#                 xindex=i
#
#             if y-word[1]-word[3]>0 and np.sqrt((x+w/2.-word[0])**2+(y+h/2.-word[1])**2)<=ydistance and np.abs(x-word[0])<im.shape[1]*0.08: #200
#                 ydistance=np.sqrt((x+w/2.-word[0])**2+(y+h/2.-word[1])**2)
#                 # absxdistance=abs(x-word[0])
#                 yindex=i
#         if xindex is not None:
#             keywordrectangles[key][word[0]].append(rectangles[xindex])
#             # print xindex
#         if yindex is not None:
#             keywordrectangles[key][word[0]].append(rectangles[yindex])


# print keywordrectangles
# translate=[]
# fig=plt.figure(0)
# axes=[]
# for i in range(20):
#     axes.append(fig.add_subplot(4,5,i+1))
# i=0
# for key in keywordrectangles.keys():
#     for key2 in keywordrectangles[key].keys():
#         for word in keywordrectangles[key][key2]:
#             x,y,w,h=word
#             # cv2.rectangle(im, (x, y), (x + w, y + h), (255,0,0), 5) ####
#             # temp=cv2.cvtColor(im[y:y+h, x:x+w,:], cv2.COLOR_BGR2GRAY)
#             temp=gray[y:y+h, x:x+w]
#             mean=np.mean(temp)/1.4
#             temp[temp<=mean]=0
#             temp[temp>mean]=255
#             # axes[i].imshow(temp, cmap="gray")
#             # if i<19:
#             #     i+=1
#             # temp=Image.fromarray(temp)
#             # translate.append(pytesseract.image_to_string(temp, lang="chi_sim"))
results=[]

if len(keyWordInfoWordAndPosition["templates_ming.png"])>0:
    hight, width = im.shape[:2]
    word = keyWordInfoWordAndPosition["templates_ming.png"][0] #通过关键字名来确定出名字的位置
    temp=[]
    # print(word)
    for rectangle in rectangles:
        if rectangle[1] < word[1] and rectangle[1]+rectangle[3]>word[1]+word[3]\
                and rectangle[0]+rectangle[2] < width/2 and rectangle[0] > word[0]+word[2]: #选取符合条件的边框，每一个边框和关键字的位置都为x,y,w,h
            temp.append(rectangle)

            # x,y,w,h=rectangle
    temp = np.array(temp)
    minx = min(temp[:,0])
    miny = min(temp[:,1])
    maxx = max(temp[:,0]+temp[:,2])
    maxy = max(temp[:,1]+temp[:,3]) #吧符合条件的边框挑出来后，去除最小的x,y,和最大x,y，从原图中截取出来关键位置

    cv2.rectangle(im, (minx, miny),(maxx, maxy), (255, 0, 0), 5)
    results.append(fm[miny:maxy, minx:maxx, :]) #把信息保存下来

if len(keyWordInfoWordAndPosition["template_jin.png"])>0:
    jinword = keyWordInfoWordAndPosition["template_jin.png"] #通过关键字 金 和费来确定收费细项的位置
    jinword = np.array(jinword)
    jinword = jinword[jinword[:,1].argsort()] #金有多个位置，挑选出位置最靠上的那一个
    jinword = jinword[0]
    # print(jinword)
    if len(keyWordInfoWordAndPosition["template_fei.png"])>0:
        rectangles = rectangles[rectangles[:,1].argsort()]
        border = rectangles[-1][1] #找出所有边框的y值的最大项，作为选取 费 字的标准
        feiword = keyWordInfoWordAndPosition["template_fei.png"]
        feiword = np.array(feiword)
        feiword = feiword[feiword[:,1] < border]
        feiword = feiword[feiword[:,1].argsort()]
        firstfeiword=feiword[0]
        lastfeiword = feiword[-1]
        save=[]
        for rectangle in rectangles:
            # if rectangle[1] > firstfeiword[1]+firstfeiword[3] and \
            #         rectangle[1]+rectangle[3] < lastfeiword[1]: #通过最考上的费和最靠下的费来确定关键内容的位置
            if rectangle[1] > jinword[1]+jinword[3] and \
                    rectangle[1]+rectangle[3] < lastfeiword[1]: #通过最考上的金和最靠下的费来确定关键内容的位置选取出边框
                save.append(rectangle)
                # x, y, w, h = rectangle
                # cv2.rectangle(im, (x, y), (x + w, y + h), (0, 255, 0), 5)
        save=np.array(save)
        h = int(np.mean(save[:,3]))
        minx = min(save[:,0])
        miny = min(save[:,1])
        maxx = max(save[:,0]+save[:,2])
        maxy = max(save[:,1]+save[:,3]) #取出边框的最小x，y和最大的x，y
        if (maxy-miny)*1./h<3.5:
            cv2.rectangle(im, (minx, miny), (maxx, maxy), (255,0,0), 10)
            results.append(fm[miny:maxy, minx:maxx, :])
        else:
            cv2.rectangle(im, (minx, miny+h), (maxx, maxy), (255,0,0), 10)
            results.append(fm[miny+h:maxy, minx:maxx, :])
        temp = []
        for rectangle in save:
            if rectangle[0] < lastfeiword[0] and \
                    rectangle[0]+rectangle[2] > lastfeiword[0]+lastfeiword[2]:
                #通过以上选出来的边框，找出倒数第二列的边框，
                # 往下挪两个h，就可以定位出总金额的位置
                temp.append(rectangle)
                x,y,w,h=rectangle
                # cv2.rectangle(im, (x, y), (x + w, y + h), (0, 255, 0), 5)
        temp = np.array(temp)
        temp = temp[temp[:,1].argsort()]
        temp = temp[-1]
        x,y,w,h=temp
        y=y+2*h
        cv2.rectangle(im, (x, y), (x + 2*w, y + h), (255,0,0), 5)
        results.append(fm[y:y+h, x:x+2*w, :])
maxx=0
maxy=0
if len(results)>0:
    results = np.array(results) #把所有的关键位置都保存到列表里，并整合成同一张图片，通过调用一次ocr来识别出所有结果
    for result in results:
        hight, width = result.shape[:2]
        maxy+=hight
        if width > maxx:
            maxx=width
    picture = np.zeros((maxy, maxx, 3), dtype = np.uint8)
    start = 0
    for result in results:
        hight, width = result.shape[:2]
        # if width < maxx:
        temp = np.zeros((hight, maxx, 3), dtype = np.uint8)
        temp[:, :width,:]=result
        picture[start:start+hight, : ,:] = temp
        start += hight
    cv2.imwrite("results.jpg", picture) #写入图片
    print(tencent_ocr("results.jpg")) #给出结果
    plt.figure(2), plt.imshow(cv2.cvtColor(picture, cv2.COLOR_BGR2RGB))

plt.figure(1), plt.imshow(cv2.cvtColor(im, cv2.COLOR_BGR2RGB))

# plt.figure(1), plt.imshow(gray, cmap="gray")
plt.show()
